Setup the environment if this is executed on Google Colab.
Make sure to change the runtime type to GPU. To do this go to Runtime -> Change runtime type -> GPU
Otherwise, rendering won't work in Google Colab.
import os
try:
import google.colab
IN_COLAB = True
except:
IN_COLAB = False
if IN_COLAB:
os.system("pip install --quiet 'x_xy[muj] @ git+https://github.com/SimiPixel/x_xy_v2'")
os.system("pip install --quiet mediapy")
import x_xy
# automatically detects colab or not
x_xy.utils.setup_colab_env()
import jax
import jax.numpy as jnp
import mediapy as media
def show_video(sys, xs, **kwargs):
media.show_video(x_xy.render(sys, [xs[i] for i in range(0, xs.shape(), 4)], camera="target", width=640, height=480, **kwargs), fps=25)
knee_xml_str = """
<x_xy model="knee_flexible_imus">
<options dt="0.01" gravity="0 0 9.81"></options>
<worldbody>
<body damping="5 5 5 25 25 25" joint="free" name="femur" pos="0.5 0.5 0.3">
<geom dim="0.1" type="xyz"></geom>
<geom dim="0.05 0.4" euler="0 90 0" mass="1" pos="0.2 0 0" type="capsule"></geom>
<body joint="frozen" name="imu1" pos="0.2 0 0.05" pos_max="0.35 0 0" pos_min="0.05 0 0">
<geom dim="0.05" type="xyz"></geom>
<geom color="orange" dim="0.05 0.05 0.02" mass="0.1" type="box"></geom>
</body>
<body damping="3" joint="ry" name="tibia" pos="0.4 0 0">
<geom dim="0.1" type="xyz"></geom>
<geom dim="0.04 0.4" euler="0 90 0" mass="1" pos="0.2 0 0" type="capsule"></geom>
<body joint="frozen" name="imu2" pos="0.2 0 0.05" pos_max="0.35 0 0" pos_min="0.05 0 0">
<geom dim="0.05" type="xyz"></geom>
<geom color="orange" dim="0.05 0.05 0.02" mass="0.1" type="box"></geom>
</body>
<geom dim="0.025 0.05 0.2" mass="0" pos="0.45 0 .1" type="box"></geom>
</body>
</body>
</worldbody>
</x_xy>
"""
sys = x_xy.load_sys_from_str(knee_xml_str)
media.show_image(x_xy.render(sys, camera="target", height=480, width=640)[0])
T = 20.0
T_initial_nomotion = 2.0
config = x_xy.MotionConfig(t_min=0.1, t_max=0.75, T=T, ang0_min=0.0, ang0_max=0.0, pos_min=-1.0, pos_max=1.0, dpos_max=0.5)
config = x_xy.join_motionconfigs([config.to_nomotion_config(), config], [T_initial_nomotion])
(X, y), (_, qs, xs, sys_mod) = x_xy.build_generator(sys, config, imu_motion_artifacts=True, dynamic_simulation=True, eager=True,
aslist=True, seed=1, sizes=1, keep_output_extras=True, imu_motion_artifacts_kwargs=dict(hide_injected_bodies=False))[0]
show_video(sys_mod, xs, show_floor=False)
sys_frozen = sys_mod.freeze("tibia").freeze("femur")
def freeze_x(q_obs):
q_frozen = jnp.concatenate(tuple(q_obs[:, sys_mod.idx_map("q")[name]] for name in ["_imu1", "imu1", "_imu2", "imu2"]), axis=-1)
return jax.vmap(lambda q: x_xy.algorithms.forward_kinematics_transforms(sys_frozen, q)[0])(q_frozen)
show_video(sys_frozen, freeze_x(qs))